///----------------------------------------------------------------------
///
/// EMC, Copyright (C) 2009 Carlos Abalde
///
/// This program is free software; you can redistribute it and/or
/// modify it under the terms of the GNU General Public License as
/// published by the Free Software Foundation; either version 2 of the
/// License, or (at your option) any later version.
///
/// This program is distributed in the hope that it will be useful,
/// but WITHOUT ANY WARRANTY; without even the implied warranty of
/// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
/// General Public License for more details.
///
/// You should have received a copy of the GNU General Public License
/// along with this program; if not, write to the Free Software
/// Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
/// 02111-1307 USA
///
///----------------------------------------------------------------------

#include <assert.h>
#include <iostream>
#include <map>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sstream>
#include <libmemcached/memcached.h>
#include "erlInterfaceUtils.h"

using namespace std;

//
// MC setup
//
memcached_st        *memc = NULL;
memcached_server_st *servers = NULL;

//
// Behaviors
//
map <string, memcached_return(*)(ETERM *)> behaviors;

////////////////////////////////////////////////////////////////////////////////
memcached_return set_boolean_behavior(ETERM *term, memcached_behavior flag)
////////////////////////////////////////////////////////////////////////////////
{
  if (ERL_IS_ATOM(term)) {
    unsigned char *value = ERL_ATOM_PTR(term);
    if (strcmp((const char *)value, "true") == 0) {
      return memcached_behavior_set(memc, flag, 1);
    } else if (strcmp((const char *)value, "false") == 0) {
      return memcached_behavior_set(memc, flag, 0);
    }
  }
  return MEMCACHED_FAILURE;
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_integer_behavior(ETERM *term, memcached_behavior flag)
////////////////////////////////////////////////////////////////////////////////
{
  if (ERL_IS_INTEGER(term)) {
    return memcached_behavior_set(memc, flag, ERL_INT_VALUE(term));
  }
  return MEMCACHED_FAILURE;
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_NO_BLOCK_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_boolean_behavior(term, MEMCACHED_BEHAVIOR_NO_BLOCK);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_TCP_NODELAY_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_boolean_behavior(term, MEMCACHED_BEHAVIOR_TCP_NODELAY);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_HASH_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  if (ERL_IS_ATOM(term)) {
    const char *value = (const char *)ERL_ATOM_PTR(term);
    if (strcmp(value, "md5") == 0) {
      return memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH , MEMCACHED_HASH_MD5);
    } else if (strcmp(value, "crc") == 0) {
      return memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH , MEMCACHED_HASH_CRC);
    } else if (strcmp(value, "fnv1_64") == 0) {
      return memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH , MEMCACHED_HASH_FNV1_64);
    } else if (strcmp(value, "fnv1A_64") == 0) {
      return memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH , MEMCACHED_HASH_FNV1A_64);
    } else if (strcmp(value, "fnv1_32") == 0) {
      return memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH , MEMCACHED_HASH_FNV1_32);
    } else if (strcmp(value, "fnv1A_32") == 0) {
      return memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH , MEMCACHED_HASH_FNV1A_32);
    }
  }
  return MEMCACHED_FAILURE;
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_DISTRIBUTION_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  if (ERL_IS_ATOM(term)) {
    const char *value = (const char *)ERL_ATOM_PTR(term);
    if (strcmp(value, "modula") == 0) {
      return memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH , MEMCACHED_DISTRIBUTION_MODULA);
    } else if (strcmp(value, "consistent") == 0) {
      return memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_HASH , MEMCACHED_DISTRIBUTION_CONSISTENT);
    }
  }
  return MEMCACHED_FAILURE;
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_CACHE_LOOKUPS_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_boolean_behavior(term, MEMCACHED_BEHAVIOR_CACHE_LOOKUPS);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_USE_UDP_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_boolean_behavior(term, MEMCACHED_BEHAVIOR_USE_UDP);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_SUPPORT_CAS_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_boolean_behavior(term, MEMCACHED_BEHAVIOR_SUPPORT_CAS);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_POLL_TIMEOUT_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_integer_behavior(term, MEMCACHED_BEHAVIOR_POLL_TIMEOUT);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_BUFFER_REQUESTS_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_boolean_behavior(term, MEMCACHED_BEHAVIOR_BUFFER_REQUESTS);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_VERIFY_KEY_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_boolean_behavior(term, MEMCACHED_BEHAVIOR_VERIFY_KEY);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_SORT_HOSTS_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_boolean_behavior(term, MEMCACHED_BEHAVIOR_SORT_HOSTS);
}

////////////////////////////////////////////////////////////////////////////////
memcached_return set_CONNECT_TIMEOUT_behavior(ETERM *term)
////////////////////////////////////////////////////////////////////////////////
{
  return set_integer_behavior(term, MEMCACHED_BEHAVIOR_CONNECT_TIMEOUT);
}

////////////////////////////////////////////////////////////////////////////////
void mc_halt()
////////////////////////////////////////////////////////////////////////////////
{
  // Deallocate memory associated with the array of memcached_server_st structures
  if (servers != NULL) {
    memcached_server_free(servers);
    servers = NULL;
  }
  // Free memcached_st structure
  if (memc != NULL) {
    memcached_free(memc);
    memc = NULL;
  }
}

////////////////////////////////////////////////////////////////////////////////
ETERM *mc_init(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  // Create a memcached_st structure
  memc = memcached_create(NULL);
  if (memc != NULL) {
    memcached_return rc;
    // Fetch servers
    ETERM *serversList = erl_element(2, argsTuple);
    while(!ERL_IS_EMPTY_LIST(serversList)) {
      // Parse tuple
      ETERM *serverTuple = ERL_CONS_HEAD(serversList);
      serversList = ERL_CONS_TAIL(serversList);
      char *serverName = erl_iolist_to_string(erl_element(1, serverTuple));
      int serverPort = ERL_INT_VALUE(erl_element(2, serverTuple));
      // Add server
      if (servers != NULL) {
	memcached_server_list_append(servers, serverName, serverPort, &rc);
      } else {
	servers = memcached_server_list_append(NULL, serverName, serverPort, &rc);
      }
      erl_free(serverName);
      if (rc != MEMCACHED_SUCCESS) {
	mc_halt();
	return erl_format("{error, ~s}", memcached_strerror(memc, rc));
      }
    }
    // Push it!
    rc = memcached_server_push(memc, servers);
    if (rc != MEMCACHED_SUCCESS) {
      mc_halt();
      return erl_format("{error, ~s}", memcached_strerror(memc, rc));
    }
    // Setup behaviors
    ETERM *behaviorList = erl_element(3, argsTuple);    
    while(!ERL_IS_EMPTY_LIST(behaviorList)) {
      // Parse tuple
      ETERM *behaviorTuple = ERL_CONS_HEAD(behaviorList);
      behaviorList = ERL_CONS_TAIL(behaviorList);
      unsigned char *behaviorName = ERL_ATOM_PTR(erl_element(1, behaviorTuple));
      ETERM *behaviorValueTerm = erl_element(2, behaviorTuple);
      stringstream ss;
      ss << behaviorName;
      // Default?
      if ((ERL_IS_ATOM(behaviorValueTerm)) &&
	  (strcmp((const char *)ERL_ATOM_PTR(behaviorValueTerm), "default") == 0)) {
	// Ignore
      } else {
	// Do it!
        if (behaviors[ss.str()] != NULL) {
	  rc = (behaviors[ss.str()])(behaviorValueTerm);
	  if (rc != MEMCACHED_SUCCESS) {
	    mc_halt();
	    return erl_format("{error, ~s}", memcached_strerror(memc, rc));
	  }
        } else {
   	  mc_halt();
	  return erl_format("{error, uknown_behavior}");
        }
      }
    }
    return erl_format("ok");
  }
  return erl_format("{error, memcached_create}");
}

////////////////////////////////////////////////////////////////////////////////
ETERM *mc_set(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  if (memc != NULL) {
    // Arguments
    char *key = erl_iolist_to_string(erl_element(2, argsTuple));
    int valueLength = ERL_INT_VALUE(erl_element(4, argsTuple));
    char *value = erl_iolist_to_string(erl_element(3, argsTuple), valueLength);
    ETERM *expirationTerm = erl_element(5, argsTuple);
    time_t expiration = (time_t)0;
    if (ERL_IS_INTEGER(expirationTerm)) {
      expiration = (time_t)ERL_INT_VALUE(expirationTerm);
    }
    // Set it!
    memcached_return rc;
    uint32_t flags = (uint32_t)0;
    rc = memcached_set(memc, key, strlen(key), value, valueLength, expiration, flags);
    erl_free(key);
    erl_free(value);
    if (rc != MEMCACHED_SUCCESS) {
      return erl_format("{error, ~s}", memcached_strerror(memc, rc));
    }
    return erl_format("ok");
  }
  return erl_format("{error, connection_not_initialized}");
}


////////////////////////////////////////////////////////////////////////////////
ETERM *mc_get(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  if (memc != NULL) {
    // Arguments
    char *key = erl_iolist_to_string(erl_element(2, argsTuple));
    // Get it!
    memcached_return rc;
    uint32_t flags = (uint32_t)0;
    size_t returnValueLength;
    char *returnValue = memcached_get(memc, key, strlen(key), &returnValueLength, &flags, &rc);
    erl_free(key);
    if (rc == MEMCACHED_SUCCESS) {
      ETERM *returnValueTerm = erl_mk_estring(returnValue, returnValueLength);
      free(returnValue);
      return erl_format("{ok, ~w}", returnValueTerm);
    } else if (rc == MEMCACHED_NOTFOUND) {
      return erl_format("{error, not_found}");
    }
    return erl_format("{error, ~s}", memcached_strerror(memc, rc));
  }
  return erl_format("{error, connection_not_initialized}");
}

////////////////////////////////////////////////////////////////////////////////
ETERM *mc_mget(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  if (memc != NULL) {
    // Arguments
    ETERM *keysList = erl_element(2, argsTuple);
    int nKeys = erl_length(keysList);
    char *keys[nKeys];
    size_t keysLength[nKeys];
    int i = 0;
    while(!ERL_IS_EMPTY_LIST(keysList)) {
      char *key = erl_iolist_to_string(ERL_CONS_HEAD(keysList));
      keys[i] = key;
      keysLength[i] = strlen(key);
      keysList = ERL_CONS_TAIL(keysList);
      i++;
    }
    // MGet it!
    memcached_return rc;
    uint32_t flags = (uint32_t)0;
    rc = memcached_mget(memc, keys, keysLength, nKeys);
    for(i=0; i<nKeys; i++) {
      erl_free(keys[i]);
    }
    if (rc == MEMCACHED_SUCCESS) {
      char returnKey[MEMCACHED_MAX_KEY];
      char *returnValue;
      size_t returnKeyLength, returnValueLength;
      ETERM *list = erl_mk_empty_list();
      while ((returnValue = memcached_fetch(memc, returnKey, &returnKeyLength, &returnValueLength, &flags, &rc))) {
	if (rc == MEMCACHED_SUCCESS) {
	  returnKey[returnKeyLength] = '\0';
          ETERM *returnValueTerm = erl_mk_estring(returnValue, returnValueLength);
 	  list = erl_cons(erl_format("{~s, ~w}", returnKey, returnValueTerm), list);
	  free(returnValue);
	} else {
	  erl_free_compound(list);
	  return erl_format("{error, ~s}", memcached_strerror(memc, rc));
	}
      }
      return erl_format("{ok, ~w}", list);
    }
    return erl_format("{error, ~s}", memcached_strerror(memc, rc));
  }
  return erl_format("{error, connection_not_initialized}");
}

////////////////////////////////////////////////////////////////////////////////
ETERM *mc_delete(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  if (memc != NULL) {
    // Arguments
    char *key = erl_iolist_to_string(erl_element(2, argsTuple));
    ETERM *expirationTerm = erl_element(3, argsTuple);
    time_t expiration = (time_t)0;
    if (ERL_IS_INTEGER(expirationTerm)) {
      expiration = (time_t)ERL_INT_VALUE(expirationTerm);
    }
    // Delete it!
    memcached_return rc;
    uint32_t flags = (uint32_t)0;
    rc = memcached_delete(memc, key, strlen(key), expiration);
    erl_free(key);
    if (rc == MEMCACHED_SUCCESS) {
      return erl_format("ok");
    } else if (rc == MEMCACHED_NOTFOUND) {
      return erl_format("{error, not_found}");
    }
    return erl_format("{error, ~s}", memcached_strerror(memc, rc));
  }
  return erl_format("{error, connection_not_initialized}");
}

////////////////////////////////////////////////////////////////////////////////
ETERM *mc_flush(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  if (memc != NULL) {
    // Arguments
    ETERM *expirationTerm = erl_element(2, argsTuple);
    time_t expiration = (time_t)0;
    if (ERL_IS_INTEGER(expirationTerm)) {
      expiration = (time_t)ERL_INT_VALUE(expirationTerm);
    }
    // Delete it!
    memcached_return rc;
    rc = memcached_flush(memc, expiration);
    if (rc == MEMCACHED_SUCCESS) {
      return erl_format("ok");
    }
    return erl_format("{error, ~s}", memcached_strerror(memc, rc));
  }
  return erl_format("{error, connection_not_initialized}");
}

////////////////////////////////////////////////////////////////////////////////
ETERM *mc_increment(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  if (memc != NULL) {
    // Arguments
    char *key = erl_iolist_to_string(erl_element(2, argsTuple));
    unsigned int offset = ERL_INT_VALUE(erl_element(3, argsTuple));
    // Increment it!
    memcached_return rc;
    uint64_t value;
    rc = memcached_increment(memc, key, strlen(key), offset, &value);
    erl_free(key);
    if (rc == MEMCACHED_SUCCESS) {
      return erl_format("{ok, ~i}", value);
    } else if (rc == MEMCACHED_NOTFOUND) {
      return erl_format("{error, not_found}");
    }
    return erl_format("{error, ~s}", memcached_strerror(memc, rc));
  }
  return erl_format("{error, connection_not_initialized}");
}


////////////////////////////////////////////////////////////////////////////////
ETERM *mc_decrement(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  if (memc != NULL) {
    // Arguments
    char *key = erl_iolist_to_string(erl_element(2, argsTuple));
    unsigned int offset = ERL_INT_VALUE(erl_element(3, argsTuple));
    // Increment it!
    memcached_return rc;
    uint64_t value;
    rc = memcached_decrement(memc, key, strlen(key), offset, &value);
    erl_free(key);
    if (rc == MEMCACHED_SUCCESS) {
      return erl_format("{ok, ~i}", value);
    } else if (rc == MEMCACHED_NOTFOUND) {
      return erl_format("{error, not_found}");
    }
    return erl_format("{error, ~s}", memcached_strerror(memc, rc));
  }
  return erl_format("{error, connection_not_initialized}");
}

////////////////////////////////////////////////////////////////////////////////
ETERM *mc_lib_version(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  const char *version = memcached_lib_version();
  return erl_format("{ok, ~w}", erl_mk_string(version));
}

////////////////////////////////////////////////////////////////////////////////
ETERM *mc_stats(ETERM *argsTuple)
////////////////////////////////////////////////////////////////////////////////
{
  if (memc != NULL) {
    memcached_return rc;
    memcached_stat_st *stats = memcached_stat(memc, NULL, &rc);
    if (rc == MEMCACHED_SUCCESS) {
      int nServers = memcached_server_count(memc);
      ETERM *slist = erl_mk_empty_list();
      for (int i= 0; i < nServers; i++) {
	char **keys, **ptr;
	keys = memcached_stat_get_keys(memc, &stats[i], &rc);
	if (rc == MEMCACHED_SUCCESS) {
	  ETERM *plist = erl_mk_empty_list();
	  for (ptr = keys; *ptr; ptr++) {
	    char *value = memcached_stat_get_value(memc, &stats[i], *ptr, &rc);
	    if (rc == MEMCACHED_SUCCESS) {
	      plist = erl_cons(erl_format("{~a, ~s}", *ptr, value), plist);
	      free(value);
	    } else {
	      free(keys);
	      erl_free_compound(slist);
	      erl_free_compound(plist);
	      return erl_format("{error, ~s}", memcached_strerror(memc, rc));
	    }
	  }
	  slist = erl_cons(erl_format("{~s, ~i, ~w}", memcached_server_name(memc, servers[i]), memcached_server_port(memc, servers[i]), plist), slist);
	  free(keys);
	} else {
	  erl_free_compound(slist);
	  return erl_format("{error, ~s}", memcached_strerror(memc, rc));
	}
      }
      free(stats);
      return erl_format("{ok, ~w}", slist);
    }
    return erl_format("{error, ~s}", memcached_strerror(memc, rc));
  }
  return erl_format("{error, connection_not_initialized}");
}

////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
int  main()
////////////////////////////////////////////////////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
{
  ETERM *result = NULL;
  ETERM *argsTuple;
  byte  *inBuffer, *outBuffer;
#ifdef DEBUG
  unsigned long allocated, freed;
#endif

  map <string, ETERM*(*)(ETERM *)> functions;

  /*
   * Register methods
   */
  functions["init"] = mc_init;
  functions["set"] = mc_set;
  functions["get"] = mc_get;
  functions["mget"] = mc_mget;
  functions["delete"] = mc_delete;
  functions["flush"] = mc_flush;
  functions["increment"] = mc_increment;
  functions["decrement"] = mc_decrement;
  functions["lib_version"] = mc_lib_version;
  functions["stats"] = mc_stats;

  /*
   * Register behaviors
   */
  behaviors["no_block"] = set_NO_BLOCK_behavior;
  behaviors["tcp_nodelay"] = set_TCP_NODELAY_behavior;
  behaviors["hash"] = set_HASH_behavior;
  behaviors["distribution"] = set_DISTRIBUTION_behavior;
  behaviors["cache_lookups"] = set_CACHE_LOOKUPS_behavior;
  behaviors["use_udp"] = set_USE_UDP_behavior;
  behaviors["support_cas"] = set_SUPPORT_CAS_behavior;
  behaviors["poll_timeout"] = set_POLL_TIMEOUT_behavior;
  behaviors["buffer_requests"] = set_BUFFER_REQUESTS_behavior;
  behaviors["verify_key"] = set_VERIFY_KEY_behavior;
  behaviors["sort_hosts"] = set_SORT_HOSTS_behavior;
  behaviors["connect_timeout"] = set_CONNECT_TIMEOUT_behavior;

  /*
   * Init Erlang - C interface
   */
  erl_init(NULL, 0);

  while ( (inBuffer=read_cmd()) != NULL )
    {
      /*
       * Read request
       */
      argsTuple = erl_decode(inBuffer);
      unsigned char *fname = ERL_ATOM_PTR(erl_element(1, argsTuple));
      stringstream ss;
      ss << fname;

#ifdef DEBUG
      fprintf(stderr, "# Call: %s\n", functionName);
#endif

      /*
       * Exec request
       */
      if (functions[ss.str()] != NULL) {
        result = (functions[ss.str()])(argsTuple);
      } else {
        result = erl_format("{error, invalid_request}");
      }

      /*
       * Send result
       */
      outBuffer = (byte *)malloc(sizeof(byte)*erl_term_len(result));
      erl_encode(result, outBuffer);
      write_cmd(outBuffer, erl_term_len(result));

      /*
       * Free resources
       */
      free(inBuffer);
      free(outBuffer);
      erl_free_compound(argsTuple);
      erl_free_compound(result);

#ifdef DEBUG
      erl_eterm_statistics(&allocated, &freed);
      fprintf(stderr, "# Currently allocated blocks: %ld\n", allocated);
      fprintf(stderr, "# Length of freelist: %ld\n", freed);
#endif

      erl_eterm_release();

    }

  mc_halt();
}
